from pettingzoo.butterfly import cooperative_pong_v5

env = cooperative_pong_v5.parallel_env(render_mode="rgb_array")
observations, infos = env.reset()
total_reward = dict()
for agent in env.possible_agents:
    total_reward[agent] = 0
while env.agents:
    # this is where you would insert your policy
    actions = {agent: env.action_space(agent).sample() for agent in env.agents}

    observations, rewards, terminations, truncations, infos = env.step(actions)
    for agent in env.agents:
        total_reward[agent] += rewards[agent]
env.close()
print(total_reward)
'''# 
env = double_dunk_v3.parallel_env(render_mode="rgb_array")
observations, infos = env.reset()



total_reward = dict()
for agent in env.possible_agents:
    total_reward[agent] = 0
total_steps = 0
while env.agents:
    # this is where you would insert your policy
    actions = {agent: env.action_space(agent).sample() for agent in env.agents}

    observations, rewards, terminations, truncations, infos = env.step(actions)
    for agent in env.agents:
        total_reward[agent] += rewards[agent]
    total_steps += 1
env.close()
print(total_reward, total_steps)
'''